
import itk
import vtk
from sys import argv
from math import pi

import connect2


fileNameVol = 'C:/Lib/jointtrack/branches/TestData/vert1.vtk'

class myReg():
    def __init__(self):
        
        # Initialize various objects for later use
        self.funccalls = 0
        self.volume = []
        self.optimizer = []
        self.fixedImage = []
        self.movingImage = []
        self.movingVolume = []
        self.pose = []
        self.imageMetric = []
        
        
        #Set up renderwindow
        #TODO: eventually we will need two render windows, one for interactive manipulation of the volume, the other to produce the DRR (no fixed image)
        self.ren = ren = vtk.vtkRenderer()
        self.renWin = renWin = vtk.vtkRenderWindow()
        renWin.AddRenderer(ren)
        
        #Create filter to capture an image from the render window
        self.w2if = w2if = vtk.vtkWindowToImageFilter()
        w2if.SetInput(renWin)
        
        self.cast = cast = vtk.vtkImageCast()
        cast.SetInput( w2if.GetOutput() )
        cast.SetOutputScalarTypeToFloat()
        cast.ClampOverflowOn()
        
        volumeTransform = vtk.vtkTransform()
    
    
    def InitializeVolume(self, fn=fileNameVol):
        """
        Open and setup a volume  for rendering.
        
        This has to be called before InitializeRegistration, since the DRR of 
        the volume will be the "moving image" in the registration.
        """
        print "opening volume file: ", fn        
        self.volume = vtk.vtkVolume()
        volumeReader = vtk.vtkStructuredPointsReader()
        volumeReader.SetFileName(fn)
        
        
        #TODO:Allow the user to modify/customize the transfer functions to change how the DRR is rendered!!
        #Default settings for the opacity transfer funciton,
        opacityTransferFunction = vtk.vtkPiecewiseFunction()        
        opacityTransferFunction.AddPoint(0.0, 0.0)
        opacityTransferFunction.AddPoint(70, 0.0)
        opacityTransferFunction.AddPoint(255, 0.5)

        # Create transfer mapping scalar value to color
        colorTransferFunction = vtk.vtkColorTransferFunction()
        colorTransferFunction.AddRGBPoint(0.0, 0.0, 0.0, 0.0)
        colorTransferFunction.AddRGBPoint(72, 1.0, 1.0, 1.0)
        colorTransferFunction.AddRGBPoint(255, 1.0, 1.0, 1.0)

        # The property describes how the data will look
        volumeProperty = vtk.vtkVolumeProperty()
        # Setup default properties
        volumeProperty.SetColor(colorTransferFunction)
        volumeProperty.SetScalarOpacity(opacityTransferFunction)
        volumeProperty.ShadeOn()
        volumeProperty.SetInterpolationTypeToLinear()

        volumeMapper = vtk.vtkVolumeTextureMapper3D()
        volumeMapper.SetInput(volumeReader.GetOutput())

        self.volume.SetProperty(volumeProperty)
        self.volume.SetMapper(volumeMapper)

        self.volume.SetOrigin(self.volume.GetCenter())
        
        #return volume
        
    def InitializeRegistration(self):
        """
        This function will setup the registration process.  It must be called before
        calculating the metric value.
        """
        
        self.ren.AddVolume(self.volume)
        
        #
        # Read the fixed images using filenames
        #
        fixedImageReader = itk.ImageFileReader.IF2.New()
        fixedFileName = 'C:/Lib/JointTrack/branches/bryan/output.png'

        fixedImageReader.SetFileName(  fixedFileName )
        fixedImageReader.Update()

        self.fixedImage = fixedImageReader.GetOutput() 
        
        #        
        # Convert from VTK image to ITK image
        #
        self.renWin.Render()
        self.w2if.Update()
        
        self.movingImage = connect2.vtk2itk(self.cast.GetOutput())
        
        #itk.echo(self.movingImage)

        #
        #  Instantiate the classes for the registration framework
        #
        registration    = itk.ImageRegistrationMethod.IF2IF2.New()
        self.imageMetric     = itk.MeanSquaresImageToImageMetric.IF2IF2.New()
        transform       = itk.CenteredRigid2DTransform.D.New()
        optimizer       = itk.RegularStepGradientDescentOptimizer.New()
        interpolator    = itk.LinearInterpolateImageFunction.IF2D.New()

        registration.SetOptimizer(      optimizer.GetPointer() )
        registration.SetTransform(      transform.GetPointer() )
        registration.SetInterpolator(   interpolator.GetPointer() )
        registration.SetMetric(         self.imageMetric.GetPointer() )
        registration.SetFixedImage(  self.fixedImage )
        registration.SetMovingImage( self.movingImage )
        registration.SetFixedImageRegion(  self.fixedImage.GetBufferedRegion() )


        #
        # Initial transform parameters 
        #
        transform.SetAngle( 0.0 );

        # center of the fixed image
        fixedSpacing = self.fixedImage.GetSpacing()
        fixedOrigin = self.fixedImage.GetOrigin()
        fixedSize = self.fixedImage.GetLargestPossibleRegion().GetSize()

        centerFixed = ( fixedOrigin.GetElement(0) + fixedSpacing.GetElement(0) * fixedSize.GetElement(0) / 2.0,
                        fixedOrigin.GetElement(1) + fixedSpacing.GetElement(1) * fixedSize.GetElement(1) / 2.0 )

        # center of the moving image 
        movingSpacing = self.movingImage.GetSpacing()
        movingOrigin = self.movingImage.GetOrigin()
        movingSize = self.movingImage.GetLargestPossibleRegion().GetSize()

        centerMoving = ( movingOrigin.GetElement(0) + movingSpacing.GetElement(0) * movingSize.GetElement(0) / 2.0,
                         movingOrigin.GetElement(1) + movingSpacing.GetElement(1) * movingSize.GetElement(1) / 2.0  )

        # transform center
        center = transform.GetCenter()
        center.SetElement( 0, centerFixed[0] )
        center.SetElement( 1, centerFixed[1] )

        # transform translation
        translation = transform.GetTranslation()
        translation.SetElement( 0, centerMoving[0] - centerFixed[0] )
        translation.SetElement( 1, centerMoving[1] - centerFixed[1] )

        initialParameters = transform.GetParameters()

        print "Initial Parameters: "
        print "Angle: %f" % (initialParameters.GetElement(0), )
        print "Center: %f, %f" % ( initialParameters.GetElement(1), initialParameters.GetElement(2) )
        print "Translation: %f, %f" % (initialParameters.GetElement(3), initialParameters.GetElement(4))

        registration.SetInitialTransformParameters( initialParameters )

        #
        # Define optimizer parameters
        #

        # optimizer scale
        translationScale = 1.0 / 100.0

        optimizerScales = itk.Array.D( transform.GetNumberOfParameters() )
        optimizerScales.SetElement(0, 1.0)
        optimizerScales.SetElement(1, translationScale/2)
        optimizerScales.SetElement(2, translationScale/2)
        optimizerScales.SetElement(3, translationScale)
        optimizerScales.SetElement(4, translationScale)

        optimizer.SetScales( optimizerScales )
        optimizer.SetMaximumStepLength( 0.2 )
        optimizer.SetMinimumStepLength( 0.002 )
        optimizer.SetNumberOfIterations( 1 )

        #
        # Start the registration process
        #

        registration.StartRegistration()


    def GetMetricValue(self, pose):
        """
        """
        #self.funccalls += 1
        metric = self.imageMetric.GetValue(pose)
        return metric
    
    def GetMetricValue1(self, ang):
        """
        """
        #self.funccalls += 1
        pose=(ang[0],256,256,0,0)
        print "pose: ", pose
        metric = self.imageMetric.GetValue(pose)
        return metric
    
    def GetMetricValue3(self,pose):
        """
        This method will trigger a DRR to be generated at the specified pose, which will then be 
        passed to the metric for a similarity measure.  
        
        The transformation will be in 3D space, but the similarity metric will be in 2D space.
        
        specified as (position(3), orientation(3))
        """
        self.SetPose(pose)
        # The image registration metric uses an internal 2D transformation.  Since we only want to
        # pass along the 2D projection as it is rendered from the DRR function, 
        # this pose will always be the same (that is, not transformed)
        #   The parameters are [angle, center_X, center_Y, translation_X, translation_Y]
        staticImagePose = (0,self.movingImageSize[0]/2,self.movingImageSize[1]/2,0,0)
        
        metric = self.imageMetric.GetValue(staticImagePose)
        
        
    
    def GenerateProjection(self,pose):
        """
        Return an 2D DRR projection of a 3D volume
        """
        #TODO: !!Allow the user to specify a transform to be applied to the volume!!
        self.renWin.render()
        
        w2if.input = self.renWin()
        ex = tvtk.PNGWriter(input=w2if.output,file_name=fn)
        w2if.update()
        ex.write() 
        
    def SetPose(self,pose):
        """
        Set the pose of the volume (which must be setup before calling this function).
        specified as (position(3), orientation(3))
        This method uses the SetPosition and SetOrientation methods of the volume object in VTK,
        
        From VTK docs:
            "Sets the orientation of the Prop3D. Orientation is 
            specified as X,Y and Z rotations in that order, but 
            they are performed as RotateZ, RotateX, and finally RotateY."

        """
        self.volume.SetPosition((pose[0],pose[1],pose[2]))
        self.volume.SetOrientation((pose[3],pose[4],pose[5]))
        

    def CB(self,pose):
        """
        Callback function for optimizer
        """
        self.funccalls += 1    
        print "I: %f, M: %f, P: %s" % (self.funccalls, self.imageMetric.GetValue(pose), pose)
        
    def SetMovingImage(self,img):
        """
        Set the moving image for registration.  This should be an ITK image.
        This method should not be used during 2D/3D registration, because the moving image will be 
        set from the projection of the 3D volume.
        """
        self.movingImageSize = []        
        self.movingImageSize.append(img.GetRequestedRegion().GetSize().GetElement(0))
        self.movingImageSize.append(img.GetRequestedRegion().GetSize().GetElement(0))
        
        self.movingImage=img
        self.metric.SetMovingImage( self.movingImage )
        
    def SetFixedImage(self,img):
        """
        Set the moving image for registration.  This should be an ITK image.  It will generally be the 
        fluoroscopic image that is being shape matched.  
        """
        
        self.fixedImageSize = []        
        self.fixedImageSize.append(img.GetRequestedRegion().GetSize().GetElement(0))
        self.fixedImageSize.append(img.GetRequestedRegion().GetSize().GetElement(0))
        
        self.fixedImage=img
        self.metric.SetFixedImage (self.fixedImage )
        
    def SetMovingVolume(self,volume):
        """
        Specify the volume to use for the registraion.
        
        """
        self.volume = volume
    
    
    def SetOptimizer(self,optimizer):
        """
        Specify the Python optimizer to use for 2D-3D registration. 
        
        Options include algorithms from the scipy.optimize collection or potentially from scikits.openopt
        """
        self.optimizer = optimizer
        

    def GetOptimizer(self):
        return self.optimizer
        
    def setImageMetric(self,metric):
        """
        Specify the ITK image 
        """
        return self.imageMetric
        
    def GetMovingVolume(self):
        return self.volume

    def GetMovingImage(self):
        return self.movingImage
        
    def GetFixedImage(self):
        return self.fixedImage
    
    def GetOptimizer(self):
        return self.optimizer
    
    def GetPose(self):
        return self.pose
    
        